from langchain_core.chat_history import InMemoryChatMessageHistory
from loguru import logger

from a0base.base_llm import pop_llm

# 创建内存聊天历史记录实例，用于存储对话消息
history = InMemoryChatMessageHistory()

# 添加用户消息到聊天历史记录
history.add_user_message("我是小明，你好")

# 调用语言模型处理聊天历史中的消息
ai_message = pop_llm.invoke(history.messages)

# 记录并输出AI回复的内容
logger.info("\n" + ai_message.content)

# 将AI回复添加到聊天历史记录中
history.add_message(ai_message)

# 添加新的用户消息到聊天历史记录
history.add_user_message("我是谁？")

# 再次调用语言模型处理更新后的聊天历史
ai_message2 = pop_llm.invoke(history.messages)

# 记录并输出第二次AI回复的内容
logger.info(ai_message2.content)

# 将第二次AI回复添加到聊天历史记录中
history.add_message(ai_message2)

# 遍历并输出所有聊天历史记录中的消息内容
logger.info("-----------")
for message in history.messages:
    logger.info(message.content)
